// http://www.gnu.org/copyleft/gpl.html
//// Code by gynvael.coldwind//vx
// !!!Use at your own risk!!!
#include<cstdio>
#include<cstring>
#include<cstdlib>
#include<windows.h>
#include<winnt.h>

using namespace std;

bool LibraryExists(const char *FileName);

int
main(void)
{
  puts("CMD with Ansi Color supported.");
  //puts("cmd.exe Ansi Patch ver.0.00d by gynvael.coldwind//vx (2009)");

  // ----
  printf("%-50s", "Checking DllSpoof.dll...");
  if(!LibraryExists("DllSpoof.dll"))
  {
    puts("ERROR\nPlease acquire or compile DllSpoof.dll first.");
    return 1;
  }
  puts("OK");

  // ----
  printf("%-50s", "Checking AnsiSupport.dll...");
  if(!LibraryExists("AnsiSupport.dll"))
  {
    puts("ERROR\nPlease acquire or compile AnsiSupport.dll first.");
    return 2;
  }
  puts("OK");

  // ----
  printf("%-50s", "Getting system directory...");
  char SystemPath[256];
  GetSystemDirectory(SystemPath, sizeof(SystemPath));
  puts(SystemPath);

  // ----
  printf("%-50s", "Acquiring a copy of cmd.exe...");
  char CmdPath[256];
  snprintf(CmdPath, sizeof(CmdPath), "%s\\cmd.exe", SystemPath);
  if(!CopyFile(CmdPath, ".\\_cmd.exe", FALSE))
//  if(!CopyFile(".\\cmd.exe", ".\\_cmd.exe", FALSE))
  {
    puts("ERROR\nCopy failed.");
    return 3;
  }
  puts("OK");

  // ----
  printf("%-50s", "Opening _cmd.exe...");
  FILE *f = fopen(".\\_cmd.exe", "r+b");
  if(!f)
  {
    puts("ERROR\nOpen failed.");
    return 4;
  }
  puts("OK");

  // ----
  printf("%-50s", "Getting OEP...");
  IMAGE_DOS_HEADER DosHeader;
  size_t Ret;
  Ret = fread(&DosHeader, 1, sizeof(DosHeader), f);
  if(Ret != sizeof(DosHeader))
  {
    puts("ERROR\nCould not read DOS header.");
    return 5;
  }

  fseek(f, DosHeader.e_lfanew, SEEK_SET);

  IMAGE_NT_HEADERS NtHeaders;
  Ret = fread(&NtHeaders, 1, sizeof(NtHeaders), f);
  if(Ret != sizeof(NtHeaders))
  {
    puts("ERROR\nCould not read NT header.");
    return 6;
  }

  DWORD OriginalEntryPoint = NtHeaders.OptionalHeader.AddressOfEntryPoint;

  printf("%.8X\n", (unsigned int)OriginalEntryPoint);

  // ----
  printf("%-50s", "Adding IMAGE_FILE_RELOCS_STRIPPED flag...");
  NtHeaders.FileHeader.Characteristics |= IMAGE_FILE_RELOCS_STRIPPED; // it's written later

  puts("OK");

  // ----
  printf("%-50s", "Resizeing last section and changing EP...");
  WORD NumberOfSections = NtHeaders.FileHeader.NumberOfSections;
  DWORD LastSectionOffset = DosHeader.e_lfanew + sizeof(NtHeaders) + sizeof(IMAGE_SECTION_HEADER) * (NumberOfSections - 1);
  DWORD NewImageSize = NtHeaders.OptionalHeader.SizeOfImage + 0x1000;
  fseek(f, LastSectionOffset, SEEK_SET);

  IMAGE_SECTION_HEADER LastSection;
  Ret = fread(&LastSection, 1, sizeof(LastSection), f);
  if(Ret != sizeof(LastSection))
  {
    puts("ERROR\nCould not read section header.");
    return 8;
  }

  DWORD NewRawDataSize = LastSection.SizeOfRawData + 0x1000;
  DWORD NewEP = LastSection.VirtualAddress + LastSection.SizeOfRawData;

  NtHeaders.OptionalHeader.SizeOfImage         = NewImageSize;
  NtHeaders.OptionalHeader.AddressOfEntryPoint = NewEP;
  LastSection.SizeOfRawData                    = NewRawDataSize;
  LastSection.Misc.VirtualSize                 = NewRawDataSize;
  LastSection.Characteristics                 |= IMAGE_SCN_CNT_CODE | 
                                                 IMAGE_SCN_MEM_EXECUTE |
						 IMAGE_SCN_MEM_WRITE |
						 IMAGE_SCN_MEM_READ;
  
  fseek(f, DosHeader.e_lfanew, SEEK_SET);
  fwrite(&NtHeaders, 1, sizeof(NtHeaders), f);
  fseek(f, LastSectionOffset, SEEK_SET);
  fwrite(&LastSection, 1, sizeof(LastSection), f);

  puts("OK");

  // ----
  printf("%-50s", "Getting address of LoadLibraryA in IAT...");
  DWORD LoadLibAddrA = 0;
  DWORD LoadLibAddrW = 0;

  DWORD IATSize = NtHeaders.OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].Size;
  DWORD IATAddr = NtHeaders.OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress;
  IMAGE_SECTION_HEADER *IATSection = NULL;

  // Find section of the IAT
  IMAGE_SECTION_HEADER *Sections = new IMAGE_SECTION_HEADER[NtHeaders.FileHeader.NumberOfSections];
  fseek(f, DosHeader.e_lfanew + sizeof(NtHeaders), SEEK_SET);
  fread(Sections, 1, sizeof(IMAGE_SECTION_HEADER) * NtHeaders.FileHeader.NumberOfSections, f);
  int i = 0;
  
  for(;i < NtHeaders.FileHeader.NumberOfSections; i++)
  {
    if(IATAddr >= Sections[i].VirtualAddress &&
       IATAddr < Sections[i].VirtualAddress + Sections[i].Misc.VirtualSize)
    {
      IATSection = &Sections[i];
      break;
    }
  }

  if(IATSection == NULL)
  {
    puts("ERROR\nCould not find IAT Section in cmd.exe.");
    return 8;
  }

  // Load IAT into memory
  BYTE *IATSectionData = new BYTE[IATSection->SizeOfRawData];
  fseek(f, IATSection->PointerToRawData, SEEK_SET);
        fread(IATSectionData, 1, IATSection->SizeOfRawData, f);

  // Scan IAT for kernel32.dll
  IMAGE_IMPORT_DESCRIPTOR *Kernel32ImpDesc = NULL;
  IMAGE_IMPORT_DESCRIPTOR *ImportDescs = (IMAGE_IMPORT_DESCRIPTOR*)&IATSectionData[IATAddr - IATSection->VirtualAddress];
  DWORD ImportDescsCounter = 0;
  while(1)
  {
    if(ImportDescsCounter * sizeof(IMAGE_IMPORT_DESCRIPTOR) >= IATSize)
      break;

    if(ImportDescs[ImportDescsCounter].Name == 0)
      break;        

    // Is the name in the same section?
    if(ImportDescs[ImportDescsCounter].Name < IATSection->VirtualAddress ||
       ImportDescs[ImportDescsCounter].Name >= IATSection->VirtualAddress + IATSection->Misc.VirtualSize)
    {
      puts("ERROR\nIAT table format not supported (are you sure cmd.exe isn't patched?).");
      return 9;
    }

    char *Name = (char*)&IATSectionData[ImportDescs[ImportDescsCounter].Name - IATSection->VirtualAddress];
    if(strnicmp(Name, "kernel32", 8) == 0)
    {
      Kernel32ImpDesc = &ImportDescs[ImportDescsCounter];
      break;
    }

    ImportDescsCounter++;
  }

  if(Kernel32ImpDesc == NULL)
  {
    puts("ERROR\nCould not find kernel32 in IAT.");
    return 9;
  }

  // Now scan the OriginalFirstThunk for the name
  DWORD ThunkCounter = 0;
  IMAGE_THUNK_DATA *ThunkData;
  if(Kernel32ImpDesc->OriginalFirstThunk < IATSection->VirtualAddress &&
     Kernel32ImpDesc->OriginalFirstThunk >= IATSection->VirtualAddress + IATSection->Misc.VirtualSize)
  {
    puts("ERROR\nIAT table format not supported (2).");
    return 10;
  }

  ThunkData = (IMAGE_THUNK_DATA*)&IATSectionData[Kernel32ImpDesc->OriginalFirstThunk - IATSection->VirtualAddress];
  while(1)
  {
    if(ThunkData[ThunkCounter].u1.AddressOfData == 0)
      break;

    if(ThunkData[ThunkCounter].u1.AddressOfData < IATSection->VirtualAddress &&
       ThunkData[ThunkCounter].u1.AddressOfData >= IATSection->VirtualAddress + IATSection->Misc.VirtualSize)
    {   
      puts("ERROR\nIAT table format not supported (2).");
      return 11;
    }

    char *Name = (char*)&IATSectionData[ThunkData[ThunkCounter].u1.AddressOfData - IATSection->VirtualAddress] + 2;

    // LoadLibraryA for XP, Vista
    if(strcmp(Name, "LoadLibraryA") == 0)
    {
      LoadLibAddrA = Kernel32ImpDesc->FirstThunk + ThunkCounter * 4 +  NtHeaders.OptionalHeader.ImageBase;
      break;
    }

    // LoadLibraryW for Windows 7
    if(strcmp(Name, "LoadLibraryW") == 0)
    {
      LoadLibAddrW = Kernel32ImpDesc->FirstThunk + ThunkCounter * 4 +  NtHeaders.OptionalHeader.ImageBase;
      break;
    }

    ThunkCounter++;
  }  

  if(LoadLibAddrA == 0 && LoadLibAddrW == 0)
  {
    puts("ERROR\nCould not find neither LoadLibraryA nor LoadLibraryW.");
    return 7;
  }

  if(LoadLibAddrA != 0)
    printf("%.8X\n", (unsigned int)LoadLibAddrA);
  else
    printf("%.8X\n", (unsigned int)LoadLibAddrW);

  // ----
  printf("%-50s", "Injecting AnsiSupport.dll loader...");

  /*char Patch[0x1000] = "\x60\xE8\x10\x00\x00\x00\x41\x6E\x73\x69\x53\x75\x70\x70\x6F\x72"
                       "\x74\x2E\x64\x6C\x6C\x00\xB8\xD3\xC0\x37\x13\xFF\xD0\x61\xB8\xBE"
                       "\xBA\xAD\x0B\xFF\xE0";*/


  // Construct the patch

  char Patch[0x1000];

  char PatchA[0x1000] = "\x60"                 // pusha
                       "\xE8\x10\x00\x00\x00"  // call GetEIP
                       "AnsiSupport.dll\0"     // lib name
                       "\xA1\xDE\xAD\xC0\xDE"  // mov eax, [LoadLibraryA]
                       "\xFF\xD0"              // call eax
                       "\x61"                  // popa
                       "\xB8\xBE\xBA\xAD\x0B"  // mov eax, OEP
                       "\xFF\xE0";             // jmp eax

  char PatchW[0x1000] = "\x60"                 // pusha
                       "\xE8\x20\x00\x00\x00"  // call GetEIP
                       "A\0n\0s\0i\0S\0u\0p\0p\0o\0r\0t\0.\0d\0l\0l\0\0\0"     // lib name
                       "\xA1\xDE\xAD\xC0\xDE"  // mov eax, [LoadLibraryA]
                       "\xFF\xD0"              // call eax
                       "\x61"                  // popa
                       "\xB8\xBE\xBA\xAD\x0B"  // mov eax, OEP
                       "\xFF\xE0";             // jmp eax

  DWORD *PtrLoadLibAddr = NULL;
  DWORD *PtrOEP         = NULL;

  if(LoadLibAddrA)
  {
    // ANSI
    memcpy(Patch, PatchA, sizeof(Patch));
    PtrLoadLibAddr = (DWORD*)&Patch[23];
    PtrOEP         = (DWORD*)&Patch[31];
    *PtrLoadLibAddr = LoadLibAddrA;
  }
  else
  {
    // Wide
    memcpy(Patch, PatchW, sizeof(Patch));
    PtrLoadLibAddr = (DWORD*)&Patch[23+16];
    PtrOEP         = (DWORD*)&Patch[31+16];
    *PtrLoadLibAddr = LoadLibAddrW;
  }

  *PtrOEP         = OriginalEntryPoint + NtHeaders.OptionalHeader.ImageBase;

  fseek(f, 0, SEEK_END);
  fwrite(Patch, 1, sizeof(Patch), f);
  puts("OK");

  // ----
  printf("%-50s", "Saving...");
  fclose(f);
  puts("OK");

  // ----
  puts("Done.");

  return 0;
}


bool LibraryExists(const char *FileName)
{
  HINSTANCE hDll;
  hDll = LoadLibraryEx(FileName, NULL, LOAD_LIBRARY_AS_DATAFILE);
  if(hDll == NULL)
    return false;

  FreeLibrary(hDll);
  return true; 
}


